package leetcode.code1819;

import java.util.PriorityQueue;

//4个堆... + 懒删
//
//大过年的，至于么，yue....呕
//
//扶墙去看看官方解答...
public class MKAverage {

	int M = (int) 1e5 + 1;
	int m, k, k2;
	int[] type;// 0k小，1中间，2k大
	int[] nums;// 各个位置的数字
	int p;
	int sm, div;// 中间累加和
	PriorityQueue<int[]> bkq;// int[]{val,index}
	PriorityQueue<int[]> skq;
	PriorityQueue<int[]> mkb;
	PriorityQueue<int[]> mks;

	final int SK = 0;// 0k小，1中间，2k大
	final int MK = 1;// 0k小，1中间，2k大
	final int BK = 2;// 0k小，1中间，2k大

	public MKAverage(int m, int k) {
		this.m = m;
		this.k = k;
		type = new int[M];// 0,1,2
		nums = new int[M];// nums
		p = 0;// 可用位置
		sm = 0;// 中间和
		k2 = k << 1;
		div = m - k2;
		bkq = new PriorityQueue<>((a, b) -> a[0] - b[0]);// k大，小根堆
		skq = new PriorityQueue<>((a, b) -> b[0] - a[0]);// k小，大根堆
		mks = new PriorityQueue<>((a, b) -> a[0] - b[0]);// rest小根堆
		mkb = new PriorityQueue<>((a, b) -> b[0] - a[0]);// rest大根堆
	}

	public void addElement(int num) {
		if (p < k2) {
			add2queue(p, num, bkq, BK);
			if (bkq.size() > k) {
				move(bkq, BK, skq, SK, null);
			}
		} else {
			int poll = p - m;
			if (poll < 0 || type[poll] == MK) {// 移除的是中间数字，塞到中间。
				add2queue(p, num, mkb, MK);
				add2queue(p, num, mks, MK);
				sm += num;
				if (poll >= 0)
					sm -= nums[poll];
				swap2m(skq, SK, mks, MK, mkb);// 小 -> 中
				swap2b(mkb, MK, bkq, BK, mks);// 中 -> 大
			} else if (type[poll] == BK) { // 移除的是k个最大的数字，塞到k个最大
				add2queue(p, num, bkq, BK);
				swap2b(mkb, MK, bkq, BK, mks);// 大 -> 中 这个交换从大到小做...,这三处交换顺序有区别！！！！
				swap2m(skq, SK, mks, MK, mkb);// 中 -> 小
			} else if (type[poll] == SK) { // 移除的是k个最小的数字，塞到k个最小
				add2queue(p, num, skq, SK);
				swap2m(skq, SK, mks, MK, mkb);// 小 -> 中
				swap2b(mkb, MK, bkq, BK, mks);// 中 -> 大
			}
		}
		p++;
	}

	// 中间 -> k大
	private void swap2b(PriorityQueue<int[]> q1, int t1, PriorityQueue<int[]> q2, int t2, PriorityQueue<int[]> q3) {
		int id1 = this.peekid(q1, t1);
		int id2 = this.peekid(q2, t2);
		if (nums[id1] > nums[id2]) {
			move(q2, t2, q1, t1, q3); // 这两顺序，先动小数！！！！
			move(q1, t1, q2, t2, null);
		}

	}

	// k小 -> 中间
	private void swap2m(PriorityQueue<int[]> q1, int t1, PriorityQueue<int[]> q2, int t2, PriorityQueue<int[]> q3) {
		int id1 = this.peekid(q1, t1);
		int id2 = this.peekid(q2, t2);
		if (nums[id1] > nums[id2]) {
			move(q2, t2, q1, t1, null);// 这两顺序，先动小数！！！！
			move(q1, t1, q2, t2, q3);
		}
	}

	// 向类型为addtype堆中添加一个数字
	private void add2queue(int index, int num, PriorityQueue<int[]> add, int addtype) {
		nums[index] = num;
		type[index] = addtype;
		add.add(new int[] { num, index });
	}

	// peek堆 顺便懒删
	private int peekid(PriorityQueue<int[]> q, int qt) {
		while (!q.isEmpty() && (q.peek()[1] <= p - m || type[q.peek()[1]] != qt)) {
			q.poll();
		}
		return q.isEmpty() ? -1 : q.peek()[1];
	}
	// poll堆 顺便懒删
	private int pollid(PriorityQueue<int[]> q, int qt) {
		while (!q.isEmpty() && (q.peek()[1] <= p - m || type[q.peek()[1]] != qt)) {
			q.poll();
		}
		return q.isEmpty() ? -1 : q.poll()[1];
	}
	// 从一个堆移动到另一个堆
	private void move(PriorityQueue<int[]> from, int fromtype, PriorityQueue<int[]> to, int totype,
			PriorityQueue<int[]> to2) {
		int idx = pollid(from, fromtype);
		if (idx == -1) {
			return;
		}
		int num = nums[idx];
		if (fromtype == MK) {
			sm -= num;
		}
		if (totype == MK) {
			sm += num;
		}
		type[idx] = totype;
		to.add(new int[] { num, idx });
		if (to2 != null) {
			to2.add(new int[] { num, idx });
		}
	}

	public int calculateMKAverage() {
		if (p < m) {
			return -1;
		}
		return sm / div;
	}

	public static void test1() {
		MKAverage obj = new MKAverage(3, 1);
		obj.addElement(3); // 当前元素为 [3]
		obj.addElement(1); // 当前元素为 [3,1]
		System.out.println(obj.calculateMKAverage()); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
		obj.addElement(10); // 当前元素为 [3,1,10]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [3,1,10]
		// 删除最小以及最大的 1 个元素后，容器为 [3]
		// [3] 的平均值等于 3/1 = 3 ，故返回 3
		obj.addElement(5); // 当前元素为 [3,1,10,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5,5]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [5,5,5]
		// 删除最小以及最大的 1 个元素后，容器
	}

	public static void test2() {

		MKAverage obj = new MKAverage(5, 2);
		obj.addElement(3); // 当前元素为 [3]
		obj.addElement(1); // 当前元素为 [3,1]
		System.out.println(obj.calculateMKAverage()); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
		obj.addElement(10); // 当前元素为 [3,1,10]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [3,1,10]
		// 删除最小以及最大的 1 个元素后，容器为 [3]
		// [3] 的平均值等于 3/1 = 3 ，故返回 3
		obj.addElement(5); // 当前元素为 [3,1,10,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5,5]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [5,5,5]
		obj.addElement(3); // 当前元素为 [3]
		obj.addElement(1); // 当前元素为 [3,1]
		System.out.println(obj.calculateMKAverage()); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
		obj.addElement(10); // 当前元素为 [3,1,10]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [3,1,10]
		// 删除最小以及最大的 1 个元素后，容器为 [3]
		// [3] 的平均值等于 3/1 = 3 ，故返回 3
		obj.addElement(5); // 当前元素为 [3,1,10,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5,5]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [5,5,5]

	}

	public static void main(String[] args) {
//		test1();
		test2();
	}

}
